"""
#!/usr/bin/env python
# -*- coding:utf-8 -*-
@Project : All_Designer
@File : Device_snap.py
@Author : RuanYaoHuang
@Time : 2023/3/18 14:57

"""
import time
from arena_api.__future__.save import Writer

class save_image():
    def __init__(self, device, pixelformat, limited, filepath, format, stopFlag=None):
        self.device = device
        self.pixelformat = pixelformat
        self.limited = limited
        self.filepattern = filepath + '/' + '<pixelformat>-[<height>,<width>]-<datetime:yy-MM-dd hh-mm-ss-fff>' + format
        self.stopflag = stopFlag
    def get_pixelformat(self):
        '''
        Generator function for pixelformat
        '''
        while True:
            yield self.device.nodemap.get_node("PixelFormat").value

    def get_height(self):
        '''
        Generator function for height
        '''
        while True:
            yield self.device.nodemap.get_node("Height").value

    def get_width(self):
        '''
        Generator function for width
        '''
        while True:
            yield self.device.nodemap.get_node("Width").value

    def get_and_save_images(self):
        self.writer = Writer()
        self.writer.register_tag("pixelformat", generator=self.get_pixelformat())
        self.writer.register_tag("height", generator=self.get_height())
        self.writer.register_tag("width", generator=self.get_width())
        self.writer.pattern = self.filepattern

        # Get device stream nodemap
        tl_stream_nodemap = self.device.tl_stream_nodemap

        # Enable stream auto negotiate packet size
        tl_stream_nodemap['StreamAutoNegotiatePacketSize'].value = True

        # Enable stream packet resend
        tl_stream_nodemap['StreamPacketResendEnable'].value = True

        nodemap = self.device.nodemap
        nodemap['PixelFormat'].value = self.pixelformat
        limit = ['1', '30', '60', '120', '180', '1', '10', '50',
         '100', '1000', '2000', '3000']
        sleepTime = int(limit[self.limited]) * 1000
        frames = int(limit[self.limited])

        with self.device.start_stream():
            # 时间限制
            if 0 <= self.limited < 5:
                while self.stopflag == False:
                    buffer = self.device.get_buffer()
                    print(f'\t\tbuffer received   | '
                          f'Width = {buffer.width} pxl, '
                          f'Height = {buffer.height} pxl, '
                          f'Pixel Format = {buffer.pixel_format.name}')
                    self.writer.save(buffer)
                    self.device.requeue_buffer(buffer)
                    time.sleep(sleepTime)
                    if self.stopflag == True:
                        break
            # 数量限制
            else:
                for i in range(frames):
                    buffer = self.device.get_buffer()
                    print(f'\t\tbuffer received   | '
                          f'Width = {buffer.width} pxl, '
                          f'Height = {buffer.height} pxl, '
                          f'Pixel Format = {buffer.pixel_format.name}')
                    print(f"Save image {i}")
                    self.writer.save(buffer)
                    self.device.requeue_buffer(buffer)

        self.device.stop_stream()